import os.path as osp
from ..utils.get_config import get_cfg
# read config file
cfg = get_cfg(osp.join(osp.abspath("./"), 'configs/config.yaml'))
DEVICE = cfg['DEVICE']


def get_loss(loader, opt, model, is_validation=True):
    # return the accuracy on a dataset

    total_loss = 0
    for data in loader:
        opt.zero_grad()
        data = data.to(DEVICE)
        embedding, pred = model(data)
        label = data.y
        if model.task == 'node':
            # train on the validation of test dataset
            mask = data.val_mask if is_validation else data.test_mask
            pred = pred[mask]
            label = label[mask]
        loss = model.loss(pred, label)
        loss.backward()
        opt.step()
        total_loss += loss.item() * data.num_graphs

    total_loss /= len(loader.dataset)
    return total_loss
